#!python
# coding=utf-8
'''
FilePath     : /code/src/eval_utils/eval.py
Description  :  
Author       : desyang
Date         : 2025-09-11 12:04:14
LastEditors  : desyang
LastEditTime : 2025-09-26 20:29:48
'''
import os
import cv2
import numpy as np

from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim
from eval_utils.glob_utils import glob_file
from multiprocessing import Pool, cpu_count

def calculate_psnr(img1, img2, max_val=255.0):
    """
    计算两幅图像的PSNR
    img1, img2: numpy数组，建议为 float 或 uint8
    max_val: 像素最大值，通常为255（8位图像）
    """
    return psnr(img1, img2, data_range=max_val)

def calculate_ssim(img1, img2, multichannel=True, win_size=7, channel_axis=None):
    """
    计算两幅图像的SSIM
    img1, img2: numpy数组
    multichannel: 是否为多通道图像（如RGB）
    win_size: 滑动窗口大小，必须是奇数且不大于图像最小边
    channel_axis: 指定通道轴（如2表示最后一维是通道），新版skimage推荐用此参数替代multichannel
    """
    # 对于新版 skimage (>=0.19)，推荐使用 channel_axis
    if channel_axis is None:
        if multichannel and img1.ndim == 3:
            channel_axis = -1  # 最后一维是通道
    
    return ssim(img1, img2, 
                win_size=win_size, 
                channel_axis=channel_axis,
                data_range=255.0)  # 固定为 uint8 图像的范围

def _calculate_single_image_pair(args):
    """
    计算单对图像的PSNR和SSIM值
    """
    file, img_dir2 = args
    filename = os.path.basename(file)
    img2_path = os.path.join(img_dir2, filename)
    
    if not os.path.exists(img2_path):
        raise FileNotFoundError(f'Image {img2_path} not found in directory')
        
    img1 = cv2.imread(file)
    img2 = cv2.imread(img2_path)

    if np.array_equal(img1, img2):
        return 100.0, 1.0
    
    if img1 is None or img2 is None:
        raise ValueError(f"Could not read image: {file} or {img2_path}")
    
    psnr_val = calculate_psnr(img1, img2)
    ssim_val = calculate_ssim(img1, img2)
    
    return psnr_val, ssim_val

def calculate_psnr_ssim(img_dir1, img_dir2, num_workers=None):
    """
    计算PSNR和SSIM，支持多进程加速

    img_dir1, img_dir2: 需要比较的图像目录  
    num_workers: 使用的CPU核心数，None表示使用所有核心

    Return: `(PSNR, SSIM)`
    """
    if num_workers is None:
        num_workers = cpu_count()
    
    files = glob_file(img_dir1)
    
    if len(files) == 0:
        raise ValueError("No images found in directory")
    
    # 准备参数列表
    args_list = [(file, img_dir2) for file in files]
    
    # 使用多进程计算
    with Pool(processes=num_workers) as pool:
        results = pool.map(_calculate_single_image_pair, args_list)
    
    # 计算平均值
    total_psnr = sum(result[0] for result in results)
    total_ssim = sum(result[1] for result in results)
    
    return total_psnr / len(files), total_ssim / len(files)


# 示例用法
if __name__ == "__main__":
    img1 = cv2.imread('images/1.jpg')
    img2 = cv2.imread('images/1.jpg')

    # 确保图像是 float 或 uint8，且范围一致
    # SSIM 支持彩色图像，会返回一个标量（各通道平均）
    ssim_value = calculate_ssim(img1, img2)
    print(f"SSIM: {ssim_value:.4f}")

    # 如果是彩色图像，建议转换为Y通道（更符合人眼）或直接计算RGB平均
    # 或者直接用RGB计算（skimage支持多通道，会自动平均）
    psnr_value = calculate_psnr(img1, img2)
    print(f"PSNR: {psnr_value:.4f} dB")